import pandas as pd
import matplotlib.pyplot as plt

pd.set_option('display.max_columns', None)  # 显示所有列，无省略
pd.set_option('display.max_rows', None)     # 显示所有行
pd.set_option('display.max_colwidth', None) # 当列内容过长时也完整显示
pd.set_option('display.width', 2000)        # 设定输出窗口宽度，防止换行断行

train_data = pd.read_csv('data/train.csv')

from sklearn.tree import DecisionTreeClassifier, plot_tree
# 假设已拆分 X_train, y_train，features 为特征名列表
clf = DecisionTreeClassifier(random_state=42, criterion='gini')

y = train_data['Survived']
train_data['Sex'] = train_data['Sex'].map({'male': 1, 'female': 0})
train_data.drop(['Survived', 'PassengerId', 'Name', 'Ticket', 'Fare', 'Cabin', 'Embarked'], axis=1, inplace=True) #删除列方法
train_data['Age'] = train_data['Age'].fillna(train_data['Age'].mean())
X = train_data

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

clf.fit(X_train, y_train)

y_predict = clf.predict(X_test)
accuracy = clf.score(X_test, y_test)
print(accuracy)

from sklearn.metrics import classification_report
print(classification_report(y_predict, y_test, target_names=['died', 'Survived']))

plt.figure(figsize=(30, 20))
plot_tree(clf, max_depth=3 ,filled=True, feature_names=X.columns, class_names=['died', 'Survived'])
plt.show()